require 'msfrpc-client'
require 'metasploit/constants'
require 'metasploit/exploit_run_description'

module Metasploit
  module Exploit
    def Exploit.start(options)

      run_details = ExploitRunDescription.new(options)
      run_details.verify

      rpc_client = get_new_metasploit_rpc_connection(run_details)

      create_workspace(rpc_client, run_details.workspace_name)

      do_nexpose_import(rpc_client, run_details)

      do_metasploit_scan(rpc_client, run_details)

      do_metasploit_audit(rpc_client, run_details)

      do_metasploit_exploit(rpc_client, run_details)

      do_metasploit_report(rpc_client, run_details)
    end

    private
    def self.get_new_metasploit_rpc_connection(run_details)
      client = Msf::RPC::Client.new(run_details.get_options)
      puts CONSTANTS::SUCCESSFUL_CONNECTION_MESSAGE

      client
    end

    def self.do_nexpose_import(rpc_client, run_details)

      if run_details.nexpose_console_name.nil? || run_details.nexpose_console_name.empty?
        puts CONSTANTS::SKIPPING_IMPORT_MESSAGE
      else
        import = rpc_client.call('pro.start_import', {'workspace' => run_details.workspace_name, 'DS_NEXPOSE_CONSOLE' => run_details.nexpose_console_name, 'DS_NEXPOSE_SITE' => run_details.workspace_name})
        wait_for_task_to_stop_running(rpc_client, CONSTANTS::IMPORTING_DATA_MESSAGE, import['task_id'])
      end
    end

    def self.create_workspace(rpc_client, workspace_name)
      rpc_client.call('pro.workspace_add', {'name' => workspace_name})
    end

    def self.do_metasploit_scan(rpc_client, run_details)
      scan = rpc_client.call('pro.start_webscan', {'workspace' => run_details.workspace_name, 'DS_URLS' => run_details.device_ip_to_scan})

      wait_for_task_to_stop_running(rpc_client, CONSTANTS::SCANNING_MESSAGE, scan['task_id'])
    end

    def self.do_metasploit_audit(rpc_client, run_details)
      audit = rpc_client.call('pro.start_webaudit', run_details.get_audit_options)

      wait_for_task_to_stop_running(rpc_client, CONSTANTS::AUDIT_MESSAGE, audit['task_id'])
    end

    def self.do_metasploit_exploit(rpc_client, run_details)
      if run_details.use_os_filter
        puts CONSTANTS::USING_OS_FILTER_MESSAGE
      else
        puts CONSTANTS::NO_OS_FILTER
      end

      if run_details.module_filter.nil? || run_details.module_filter.empty?
        puts CONSTANTS::NO_MODULE_FILTER
      else
        puts CONSTANTS::USING_MODULE_FILTER
        puts run_details.module_filter
      end

      sploit = rpc_client.call('pro.start_exploit', run_details.get_exploit_options)

      wait_for_task_to_stop_running(rpc_client, CONSTANTS::EXPLOIT_MESSAGE, sploit['task_id'])
    end

    def self.do_metasploit_report(rpc_client, run_details)
      report_type = run_details.report_type
      if report_type.nil? || report_type.empty?
        puts CONSTANTS::NO_REPORT_TYPE_MESSAGE
      else
        puts "Generating a #{report_type} Report"

        report = rpc_client.call('pro.start_report', run_details.get_report_options)

        self.write_report_after_generation(rpc_client, run_details, report['report_id'])
      end
    end

    def self.wait_for_task_to_stop_running(rpc_client, status_message, task_id)
      loop do
        sleep(3)
        status = rpc_client.call('pro.task_status', task_id)
        puts status_message
        progress = status.fetch(task_id, {}).fetch('progress', {})
        puts "This task is #{progress} % complete"
        info = status.fetch(task_id, {}).fetch('info', {})
        puts "I am currently executing: #{info}"
        status = status.fetch(task_id, {}).fetch('status', {})

        break if status != CONSTANTS::RUNNING_IMPORT_STATUS
      end
    end

    def self.write_report_after_generation(rpc_client, run_details, report_id)
      (1..100).each do
        download = rpc_client.call('pro.report_download', report_id)
        if download && download['report_artifacts'] && download['report_artifacts'].length > 0
          download['report_artifacts'].each_with_index do |artifact, index|
            File.open("#{run_details.workspace_name}_#{index}#{File.extname(artifact['file_path'])}", "wb") do |fd|
              fd.write(artifact['data'])
            end
            puts "Report File #{index} Written"
          end
          break
        else
          puts "Waiting for report to generate"
          sleep(3)
        end
      end
    end
  end
end
